# train.py
import argparse
import os
import torch
import torch.distributed as dist # 新增: 用于分布式训练
from torch.nn.parallel import DistributedDataParallel as DDP # 新增: 用于分布式训练
from torch.utils.data.distributed import DistributedSampler # 新增: 用于分布式训练
from datasets import load_dataset, load_from_disk # 修改: 增加 load_from_disk
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
)

def main():
    parser = argparse.ArgumentParser(description="微调 TinyLlama 模型")
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        default="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        help="预训练模型的路径或 HuggingFace 模型名称。",
    )
    parser.add_argument(
        "--local_model_path",
        type=str,
        default=None,
        help="本地预训练模型的路径 (如果使用本地模型)。"
    )
    parser.add_argument(
        "--dataset_name",
        type=str,
        default="wikitext",
        help="HuggingFace 数据集名称。",
    )
    parser.add_argument(
        "--dataset_config_name",
        type=str,
        default="wikitext-2-raw-v1", # 一个较小的数据集，适合演示
        help="HuggingFace 数据集配置名称。",
    )
    parser.add_argument(
        "--local_dataset_path",
        type=str,
        default=None,
        help="本地数据集的路径 (如果使用本地数据集)。"
    )
    parser.add_argument(
        "--use_local_data",
        action="store_true",
        help="是否从本地路径加载模型和数据集。"
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="./model_output",
        help="训练后模型的输出目录。",
    )
    parser.add_argument(
        "--num_train_epochs",
        type=int,
        default=1, # 演示目的，只训练1个epoch
        help="训练的总轮数。",
    )
    parser.add_argument(
        "--per_device_train_batch_size",
        type=int,
        default=1, # 根据你的GPU显存调整
        help="每个设备上的训练批量大小。",
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=4, # 梯度累积步骤
        help="梯度累积的步数，用于模拟更大的批量大小。",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=5e-5,
        help="初始学习率。",
    )
    parser.add_argument(
        "--logging_steps",
        type=int,
        default=10,
        help="每 N 步记录一次日志。",
    )
    parser.add_argument(
        "--save_steps",
        type=int,
        default=500, # 每500步保存一次checkpoint
        help="每 N 步保存一次模型 checkpoint。",
    )
    parser.add_argument(
        "--max_seq_length",
        type=int,
        default=512, # 最大序列长度
        help="输入序列的最大长度。",
    )
    parser.add_argument(
        "--no_cuda",
        action="store_true",
        help="不使用 CUDA，即使可用。"
    )
    # 新增: 分布式训练参数 (通常由 torchrun/launch.py 自动设置)
    parser.add_argument(
        "--local_rank", 
        type=int, 
        default=-1, 
        help="分布式训练的本地进程排名 (由启动器自动设置)"
    )

    args = parser.parse_args()

    print(f"训练参数: {args}")

    # 初始化分布式环境
    is_distributed = args.local_rank != -1
    if is_distributed:
        if not dist.is_available():
            raise RuntimeError("Distributed training requested but Pytorch distributed is not available.")
        # NCCL is a good default backend for NVIDIA GPUs.
        # dist.init_process_group(backend='nccl', init_method='env://') # 如果使用 env://
        # Kubeflow PyTorchJob 通常会设置 MASTER_ADDR 和 MASTER_PORT
        # 如果通过 torchrun 启动，它也会处理这些环境变量
        if os.environ.get('MASTER_ADDR') and os.environ.get('MASTER_PORT') and os.environ.get('WORLD_SIZE') and os.environ.get('RANK'):
            print(f"Initializing process group via env variables: MASTER_ADDR={os.environ['MASTER_ADDR']}, MASTER_PORT={os.environ['MASTER_PORT']}")
            dist.init_process_group(backend='nccl')
            args.local_rank = int(os.environ['LOCAL_RANK']) #确保 local_rank 被正确设置
            torch.cuda.set_device(args.local_rank) # 为当前进程设置CUDA设备
            print(f"Distributed training initialized. World size: {dist.get_world_size()}, Rank: {dist.get_rank()}, Local Rank: {args.local_rank}")
        else:
            print("Warning: Distributed training environment variables (MASTER_ADDR, MASTER_PORT, etc.) not fully set. Attempting to initialize with default env://")
            try:
                dist.init_process_group(backend='nccl', init_method='env://')
                args.local_rank = int(os.environ.get('LOCAL_RANK', args.local_rank))
                if args.local_rank != -1:
                    torch.cuda.set_device(args.local_rank)
                print(f"Distributed training initialized via env://. World size: {dist.get_world_size()}, Rank: {dist.get_rank()}, Local Rank: {args.local_rank}")
            except Exception as e:
                raise RuntimeError(f"Failed to initialize process group with env://. Ensure MASTER_ADDR, MASTER_PORT, WORLD_SIZE, RANK, LOCAL_RANK are set. Error: {e}")
    else:
        print("Not a distributed run.")


    # 检查 CUDA 可用性
    use_cuda = not args.no_cuda and torch.cuda.is_available()
    if is_distributed:
        # 在分布式模式下，设备由 local_rank 决定
        device = torch.device(f"cuda:{args.local_rank}")
    else:
        device = torch.device("cuda" if use_cuda else "cpu")
    
    print(f"使用的设备: {device}")
    if use_cuda:
        print(f"可用的 GPU 数量: {torch.cuda.device_count()}")
        if not is_distributed or args.local_rank == 0:
             print(f"当前 GPU (主进程或单机): {torch.cuda.get_device_name(0 if not is_distributed else args.local_rank)}")

    # 1. 加载 Tokenizer 和模型
    model_load_path = args.local_model_path if args.use_local_data and args.local_model_path else args.model_name_or_path
    print(f"从 '{model_load_path}' 加载 Tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(model_load_path)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        print(f"Tokenizer pad_token 未设置，已设置为 eos_token: {tokenizer.eos_token}")

    print(f"从 '{model_load_path}' 加载模型...")
    model = AutoModelForCausalLM.from_pretrained(
        model_load_path,
        # trust_remote_code=True # 如果模型包含自定义代码，则需要
    )
    model.to(device) # 将模型移动到指定设备

    # 如果是分布式训练，则封装模型
    if is_distributed:
        model = DDP(model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=False)
        print(f"模型已封装为 DistributedDataParallel on rank {args.local_rank}")

    # 2. 加载和预处理数据集
    if args.use_local_data and args.local_dataset_path:
        print(f"从本地路径 '{args.local_dataset_path}' 加载数据集...")
        try:
            raw_datasets = load_from_disk(args.local_dataset_path)
            print(f"成功从本地加载数据集。Splits: {list(raw_datasets.keys())}")
        except Exception as e:
            raise ValueError(f"无法从本地路径 '{args.local_dataset_path}' 加载数据集: {e}")
    else:
        print(f"加载数据集 '{args.dataset_name}' (配置: '{args.dataset_config_name}')...")
        raw_datasets = load_dataset(args.dataset_name, args.dataset_config_name)

    # 筛选掉空文本，并选择训练集的一个小子集进行快速演示
    # 对于实际训练，你应该使用完整的数据集或更大的子集
    if "train" in raw_datasets:
        train_dataset = raw_datasets["train"].filter(lambda example: len(example["text"]) > 0)
        # train_dataset = train_dataset.select(range(min(1000, len(train_dataset)))) # 取前1000条作为演示
        print(f"处理后的训练样本数: {len(train_dataset)}")
    else:
        raise ValueError("数据集中未找到 'train' 部分。请检查数据集结构。")

    # 定义文本列名
    column_names = train_dataset.column_names
    text_column_name = "text" if "text" in column_names else column_names[0]

    def tokenize_function(examples):
        # 对文本进行 tokenize
        tokenized_output = tokenizer(
            examples[text_column_name],
            truncation=True,
            max_length=args.max_seq_length,
            padding="max_length", # 或者 "longest"
        )
        return tokenized_output

    print("对数据集进行 Tokenize...")
    tokenized_datasets = train_dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=column_names, # 移除原始文本列
        desc="运行 tokenizer 在数据集上",
    )
    print(f"Tokenized 数据集的一个样本: {tokenized_datasets[0]}")


    # 3. 设置 Data Collator
    # DataCollatorForLanguageModeling 会自动创建语言模型的标签 (labels 通常是 input_ids 的移位)
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    # 4. 设置训练参数
    # 在分布式训练中，确保只有主进程 (rank 0) 执行可能冲突的操作，如写文件
    # Hugging Face Trainer 会自动处理大部分分布式训练的逻辑
    # 我们需要确保 TrainingArguments 知道 local_rank
    print("设置训练参数...")
    training_args_dict = {
        "output_dir": args.output_dir,
        "overwrite_output_dir": True,
        "num_train_epochs": args.num_train_epochs,
        "per_device_train_batch_size": args.per_device_train_batch_size,
        "gradient_accumulation_steps": args.gradient_accumulation_steps,
        "save_steps": args.save_steps,
        "save_total_limit": 2, # 最多保存2个checkpoint
        "logging_steps": args.logging_steps,
        "learning_rate": args.learning_rate,
        "fp16": use_cuda, # 如果使用 CUDA，则启用 fp16 混合精度训练
        "report_to": "tensorboard", # 或者 "wandb", "none"
        "dataloader_num_workers": 2, # 根据你的系统调整
        "no_cuda": args.no_cuda,
        # "ddp_find_unused_parameters": False, # 如果遇到 DDP 问题可以尝试
    }

    if is_distributed:
        # Trainer 会自动处理 local_rank if DDP is used and env vars are set by torchrun
        # training_args_dict["local_rank"] = args.local_rank # 通常不需要显式设置，Trainer会检测
        # training_args_dict["ddp_find_unused_parameters"] = False # 某些情况下需要
        pass # Trainer 会自动处理分布式设置

    training_args = TrainingArguments(**training_args_dict)

    # 5. 初始化 Trainer
    print("初始化 Trainer...")
    
    # 为分布式训练创建 DistributedSampler
    # Trainer 内部会处理 DistributedSampler，如果 train_dataset 是普通 Dataset
    # 并且检测到分布式环境。但如果想更明确，可以手动创建。
    # train_sampler = DistributedSampler(tokenized_datasets) if is_distributed else None
    # Trainer 会自动处理 sampler，所以这里不需要手动创建 train_sampler

    trainer = Trainer(
        model=model, # DDP-wrapped model if distributed
        args=training_args,
        train_dataset=tokenized_datasets,
        # eval_dataset=tokenized_datasets["validation"], # 如果有验证集
        data_collator=data_collator,
        tokenizer=tokenizer,
        # sampler=train_sampler, # Trainer 会自动处理
    )

    # 6. 开始训练
    # 确保只有主进程打印某些信息或执行保存操作（Trainer 会处理大部分）
    if not is_distributed or dist.get_rank() == 0:
        print("开始训练 (主进程)...")
    else:
        print(f"开始训练 (工作进程 rank {dist.get_rank()})...")
    
    train_result = trainer.train()
    
    if not is_distributed or dist.get_rank() == 0:
        print("训练完成 (主进程)。")

    # 7. 保存模型和训练状态 (Trainer 会确保只在主进程保存)
    if not is_distributed or dist.get_rank() == 0:
        print(f"保存最终模型到 {args.output_dir} (主进程)...")
    # trainer.save_model() # Trainer.train() 结束后会自动保存 (如果 save_strategy 是 'steps' 或 'epoch')
    # 如果需要在训练结束后强制保存，可以调用。Trainer 会处理分布式保存。
    # 对于DDP模型，需要获取原始模型进行保存
    # model_to_save = model.module if hasattr(model, 'module') else model
    # model_to_save.save_pretrained(args.output_dir)
    # tokenizer.save_pretrained(args.output_dir)
    # trainer.save_state() # 保存训练状态
    # Hugging Face Trainer 在训练结束时会自动调用 save_model 和 save_state (基于 save_strategy)
    # 如果 save_strategy='no', 则需要手动调用
    # 默认 save_strategy='steps', 所以在最后一步会保存
    # 如果需要确保在训练结束后总是保存最新的模型，可以显式调用
    if trainer.is_world_process_zero(): # 确保只在主进程操作
        print(f"显式保存最终模型和状态到 {args.output_dir}...")
        trainer.save_model() 
        trainer.save_state()

    # 打印一些训练结果统计
    if trainer.is_world_process_zero(): # 确保只在主进程操作
        metrics = train_result.metrics
        metrics["train_samples"] = len(tokenized_datasets)
        trainer.log_metrics("train", metrics)
        trainer.save_metrics("train", metrics)

    if is_distributed:
        dist.destroy_process_group()
        print("分布式进程组已销毁。")

if __name__ == "__main__":
    main()
